import { Key, TreeNode } from './interface'
import { filter as presetFilter } from './preset'

/**
 * 改写自naive-ui
 * @param tree
 * @param filter
 * @param pattern
 * @copyright https://github.com/TuSimple/naive-ui/blob/main/src/tree/src/utils.ts
 */
export function filterTree<Raw>(
  tree: TreeNode<Raw>[],
  filter: (pattern: string, v: string) => boolean = presetFilter,
  pattern: string
): TreeNode<Raw>[] {
  const visitedTailKeys = new Set<Key>()
  const visitedNonTailKeys = new Set<Key>()
  const filteredTree: TreeNode<Raw>[] = []
  function visit(t: TreeNode<Raw>[]): void {
    t.forEach((n) => {
      if (filter(pattern, n.label)) {
        visitedTailKeys.add(n.key)
        for (let i = n.path.length - 2; i >= 0; --i) {
          const key: Key = n.path[i]
          if (!visitedNonTailKeys.has(key)) {
            visitedNonTailKeys.add(key)
            if (visitedTailKeys.has(key)) {
              visitedTailKeys.delete(key)
            }
          } else {
            break
          }
        }
      }
      const children = n.children
      if (children) {
        visit(children)
      }
    })
  }
  visit(tree)
  function build(t: TreeNode<Raw>[], sibs: TreeNode<Raw>[]): void {
    t.forEach((n) => {
      const key = n.key
      const isVisitedTail = visitedTailKeys.has(key)
      const isVisitedNonTail = visitedNonTailKeys.has(key)
      if (!isVisitedTail && !isVisitedNonTail) return
      const children = n.children
      if (children) {
        if (isVisitedTail) {
          // If it is visited path tail, use origin node
          sibs.push(n)
        } else {
          // It it is not visited path tail, use cloned node
          const clonedNode = { ...n, children: [] }
          sibs.push(clonedNode)
          build(children, clonedNode.children as TreeNode<Raw>[])
        }
      } else {
        sibs.push(n)
      }
    })
  }
  build(tree, filteredTree)
  return filteredTree
}
